#!/usr/bin/python

#load smiles in table, identifier in supplier table

import string,os, sys
import MySQLdb, cPickle, mysql
import gzip 
from cStringIO import StringIO

if len(sys.argv) <> 5:
	print 'mysql_get_dockable.py <prefix of database> <number of molecules per database> <query file (get wanted ids)> <query file (get stereo blobs)> '
	print 'Note: second query must end with "where statement", e. g. "... where id ="'
	sys.exit(1)

name = sys.argv[1]
chunks = int(sys.argv[2])
query_get_ids = open(sys.argv[3], 'r')
query_get_db = open(sys.argv[4], 'r')


#-------------------------------------
def store_molis(header,molis):
	#write gzip file
	out_file = gzip.GzipFile('tmp.db.gz' , 'wb')
	out_file.write(header)
	out_file.write(molis)
	out_file.close()

	#read gzip file
	
	zipped_file = open('tmp.db.gz' , 'rb')
	content = zipped_file.read()
	content = cPickle.dumps(content,2) #cPickle.dump converts data into binary structure
	command = "insert into test (dockable) values (%s)"
	cursor.execute(command, ((content)))


#-------------------------------------
def get_header(file):
	header = file.readline()
	new_line = file.readline()
	while new_line.find('(') <> -1:
		header = header + new_line
		new_line = file.readline()

	#new line is family line, append to molecule
	return header,new_line
#-------------------------------------
def format_id (id):
	id = str(id)
	id = id.rjust(8)
	id = id.replace(' ', '0')
	id = 'V' + id
	return id
#-------------------------------------
def write_entry(out_file,id,blob,old_header,first):
	#print id
	#print blob
	content = cPickle.loads(blob.tostring()) 	#cPickle.loads converts data from binary data into non-binary data
	new_file = StringIO(content) 	#converts data into a pseudo-file
	unzipped = gzip.GzipFile(fileobj=new_file, mode = 'r')

	new_header, moli = get_header(unzipped)
	if new_header <> old_header and old_header <> '':
		print 'Warning: Colouring does not match', id
	if first:
		out_file.write(new_header) # first entry
	moli = moli + unzipped.read()
	unzipped.close()
	new_file.close()
	moli = moli.replace('xxxxxxxx', id)
	#print moli
	#print id
	if moli.find('0     0     0     0') == -1:
		out_file.write(moli)
	return new_header	
#-------------------------------------
def write_results(out_file,results):
	#get first entry to check the header
	row = results[0]
	id = format_id(row[0])
	blob = row[1]
	first = True
	first_header = write_entry(out_file,id,blob,'',first)

	first = False


	for row in range(1,len(results)):
		id = format_id(results[row][0])
		blob = results[row][1]
		header = write_entry(out_file,id,blob,first_header,first)
		
	conn.close ()

#-------------------------------------

conn=mysql.connect2server('', 'webuser', 'purchasable')  			  
cursor = conn.cursor ()



command = query_get_ids.readline()  #this is the query that gets all id numbers for which a dockable db is wanted



print command
cursor.execute(command)

ids = cursor.fetchall()

print len(ids), ' id numbers found'

counter = 0.0
out_file = open(name + '_' + str(int(counter)) + '.db' , 'w')


basic_command =  string.strip(query_get_db.readline())  #this query get the stereo_blobs, has to end with id = 

#get first entry to check header
id = str(ids[0][0])
command = basic_command + id
print command
cursor.execute(command)
results=cursor.fetchall()
row = results[0]
id = format_id(row[0])
blob = row[1]
first = True
first_header = write_entry(out_file,id,blob,'',first)

first = False
#write missing entries form this id number
for entry in range(1,len(results)):
	row = results[entry]
	id = format_id(row[0])
	blob = row[1]
	header = write_entry(out_file,id,blob,first_header,first)



old_counter = 0

for entry in range(1,len(ids)):
	id = str(ids[entry][0])
	#print 'supplier id', id
	if int(counter / chunks) > old_counter:
		out_file.close()
		out_file = open(name + '_' + str(int(old_counter+1)) + '.db' , 'w')
		out_file.write(first_header)
		old_counter = int(counter / chunks) 
		
	command = basic_command + str(id)
	#print command
	cursor.execute(command)
	results=cursor.fetchall()
	for row in results:
		#print id
		id = format_id(row[0])
		#print id
		blob = row[1]
		header = write_entry(out_file,id,blob,first_header,first)
		
		counter = counter + 1
conn.close ()
out_file.close()

	
